package com.seago.iceberg.repository;

import cn.hutool.core.lang.UUID;
import com.seago.iceberg.config.IcebergInit;
import org.apache.iceberg.*;
import org.apache.iceberg.data.GenericRecord;
import org.apache.iceberg.data.IcebergGenerics;
import org.apache.iceberg.data.Record;
import org.apache.iceberg.data.parquet.GenericParquetWriter;
import org.apache.iceberg.expressions.Expressions;
import org.apache.iceberg.io.CloseableIterable;
import org.apache.iceberg.io.DataWriter;
import org.apache.iceberg.io.OutputFile;
import org.apache.iceberg.parquet.Parquet;
import org.apache.iceberg.types.Types;

import java.beans.IntrospectionException;
import java.beans.Introspector;
import java.beans.PropertyDescriptor;
import java.io.IOException;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Title:
 * <p>
 * description:
 *
 * @Company: 中国船舶集团海舟系统技术有限公司-武汉
 * @author: yangs
 * @since: 2023-05-08 09:28
 * @version: 1.0
 **/
public class SimpleIcebergRepository<T, ID> implements IcebergRepository<T, ID> {

	private Class<T> entityClass;
	private Table icebergTable;

	private Class<ID> idType;

	public SimpleIcebergRepository(Class<T> entityClass, Class<ID> idType) {
		this.entityClass = entityClass;
		this.idType = idType;
		this.icebergTable = IcebergInit.getIcebergTable(entityClass);
	}

	@Override
	public void save(T entity) throws Exception {
		Transaction t = icebergTable.newTransaction();
		GenericRecord record = toIcebergRecord(entity);
		// 2. 将记录写入parquet文件
		String filepath = icebergTable.location() + "/data/" + UUID.randomUUID() + ".parquet";
		OutputFile file = icebergTable.io().newOutputFile(filepath);
		DataWriter<GenericRecord> dataWriter =
				Parquet.writeData(file)
						.schema(icebergTable.schema())
						.createWriterFunc(GenericParquetWriter::buildWriter)
						.overwrite()
						.withSpec(PartitionSpec.unpartitioned())
						.build();
		try {
			dataWriter.write(record);
		} finally {
			try {
				dataWriter.close();
			} catch (IOException e) {
				throw new RuntimeException(e);
			}
		}

		// 3. 将文件写入table中
		DataFile dataFile = dataWriter.toDataFile();
		t.newAppend().appendFile(dataFile).commit();
		t.commitTransaction();
	}


	@Override
	public List<T> findAll() {
		List<T> result = new ArrayList<>();
		IcebergGenerics.ScanBuilder scanBuilder = IcebergGenerics.read(icebergTable);
		// 查询全部
		CloseableIterable<Record> records = scanBuilder.build();
		for (Record record : records) {
			// 将每个Record对象转换为T对象，并将其添加到结果列表中
			T obj = fromIcebergRecord(record, entityClass);
			result.add(obj);
		}
		return result;
	}

	@Override
	public List<T> getById(ID id) {
		System.out.println("getById = " + id + "," + "当前类=SimpleIcebergRepository.getById()");
		List<T> result = new ArrayList<>();
		IcebergGenerics.ScanBuilder scanBuilder = IcebergGenerics.read(icebergTable);
		scanBuilder.select();

		CloseableIterable<Record> records = scanBuilder.select("id", "name")
				.where(Expressions.equal("id", id)).build();
		for (Record record : records) {
			// 将每个Record对象转换为T对象，并将其添加到结果列表中
			T obj = fromIcebergRecord(record, entityClass);
			result.add(obj);
		}
		return result;
	}

	protected GenericRecord toIcebergRecord(T entity) {
		try {
			Schema schema = icebergTable.schema();
			GenericRecord record = GenericRecord.create(schema);
			PropertyDescriptor[] propertyDescriptors = Introspector.getBeanInfo(entityClass).getPropertyDescriptors();
			Map<String, PropertyDescriptor> descriptorMap = new HashMap<>();
			for (PropertyDescriptor propertyDescriptor : propertyDescriptors) {
				String columnName = propertyDescriptor.getName();
				descriptorMap.put(columnName, propertyDescriptor);
			}
			// 将实体类的属性值复制到Iceberg记录
			schema.columns().forEach(column -> {
				PropertyDescriptor propertyDescriptor = descriptorMap.get(column.name());
				if (propertyDescriptor != null) {
					try {
						Object value = propertyDescriptor.getReadMethod().invoke(entity);
						record.setField(column.name(), value);
					} catch (IllegalAccessException | InvocationTargetException e) {
						throw new RuntimeException("读取列的属性值时出错: " + column.name(), e);
					}
				}
			});
			return record;
		} catch (IntrospectionException e) {
			throw new RuntimeException("错误的实体类: " + entityClass.getName(), e);
		}
	}

	private T fromIcebergRecord(Record record, Class<T> clazz) {
		T obj;
		try {
			obj = clazz.getDeclaredConstructor().newInstance();
			Types.StructType structType = record.struct();
			for (Types.NestedField field : structType.fields()) {
				String fieldName = field.name();
				Object fieldValue = record.getField(fieldName);
				Field objField = clazz.getDeclaredField(fieldName);
				objField.setAccessible(true);
				objField.set(obj, fieldValue);
			}
		} catch (Exception e) {
			throw new RuntimeException("将Iceberg record转换为类型对象时出错 " + clazz.getName(), e);
		}
		return obj;
	}
}
